{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bd2f8165",
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'backtrader'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Input \u001b[1;32mIn [4]\u001b[0m, in \u001b[0;36m<cell line: 2>\u001b[1;34m()\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtushare\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mts\u001b[39;00m\n\u001b[1;32m----> 2\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mbacktrader\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mbt\u001b[39;00m\n\u001b[0;32m      3\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mpandas\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mpd\u001b[39;00m\n\u001b[0;32m      4\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mnumpy\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnp\u001b[39;00m\n",
      "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'backtrader'"
     ]
    }
   ],
   "source": [
    "import tushare as ts\n",
    "import backtrader as bt\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# 初始化tushare\n",
    "ts.set_token('6f4ed3cc824c2d45087664dec52d197dfc411c36f4ed520ed8c35c50')\n",
    "pro = ts.pro_api()\n",
    "\n",
    "# 选股与排序\n",
    "def select_and_sort_stocks():\n",
    "    # 选择市值较大且市盈率较低的股票\n",
    "    stocks = pro.stock_basic(exchange='', list_status='L', fields='ts_code,total_mv,pe')\n",
    "    stocks = stocks[(stocks['total_mv'] > 1000) & (stocks['pe'] > 0)].sort_values(by=['total_mv', 'pe'],\n",
    "                                                                                  ascending=[False, True])\n",
    "    return stocks['ts_code'].head(5).tolist()\n",
    "\n",
    "# 交易信号生成策略（双均线策略）\n",
    "class DualMovingAverageStrategy(bt.Strategy):\n",
    "    params = (\n",
    "        ('short_window', 5),\n",
    "        ('long_window', 20),\n",
    "        ('stop_loss', 0.1),\n",
    "        ('take_profit', 0.2)\n",
    "    )\n",
    "\n",
    "    def __init__(self):\n",
    "        self.short_ma = bt.indicators.SimpleMovingAverage(\n",
    "            self.data.close, period=self.params.short_window)\n",
    "        self.long_ma = bt.indicators.SimpleMovingAverage(\n",
    "            self.data.close, period=self.params.long_window)\n",
    "        self.crossover = bt.indicators.CrossOver(self.short_ma, self.long_ma)\n",
    "        self.buy_price = None\n",
    "\n",
    "    def next(self):\n",
    "        if not self.position:\n",
    "            if self.crossover > 0:\n",
    "                self.buy()\n",
    "                self.buy_price = self.data.close[0]\n",
    "        else:\n",
    "            if self.data.close[0] <= self.buy_price * (1 - self.params.stop_loss):\n",
    "                self.sell()\n",
    "                self.buy_price = None\n",
    "            elif self.data.close[0] >= self.buy_price * (1 + self.params.take_profit):\n",
    "                self.sell()\n",
    "                self.buy_price = None\n",
    "\n",
    "# 回测函数\n",
    "def backtest(selected_stocks, time_scale='D'):\n",
    "    cerebro = bt.Cerebro()\n",
    "    all_returns = []\n",
    "\n",
    "    for stock in selected_stocks:\n",
    "        if time_scale == 'D':\n",
    "            df = pro.daily(ts_code=stock, start_date='20200101', end_date='20231231')\n",
    "        elif time_scale == 'W':\n",
    "            df = pro.weekly(ts_code=stock, start_date='20200101', end_date='20231231')\n",
    "        elif time_scale == '1min':\n",
    "            df = pro.mins(ts_code=stock, freq='1min', start_date='20230101', end_date='20231231')\n",
    "        else:\n",
    "            raise ValueError(\"不支持的时间尺度\")\n",
    "\n",
    "        df['trade_date'] = pd.to_datetime(df['trade_date'])\n",
    "        df.set_index('trade_date', inplace=True)\n",
    "        df = df[['open', 'high', 'low', 'close', 'vol']]\n",
    "        df.rename(columns={'vol': 'volume'}, inplace=True)\n",
    "\n",
    "        data = bt.feeds.PandasData(dataname=df)\n",
    "        cerebro.adddata(data)\n",
    "\n",
    "    cerebro.addstrategy(DualMovingAverageStrategy)\n",
    "    cerebro.broker.setcash(100000.0)\n",
    "    cerebro.broker.setcommission(commission=0.001)\n",
    "\n",
    "    print('初始资金: %.2f' % cerebro.broker.getvalue())\n",
    "    results = cerebro.run()\n",
    "    print('最终资金: %.2f' % cerebro.broker.getvalue())\n",
    "\n",
    "    for i, result in enumerate(results):\n",
    "        portfolio_value = [value[0] for value in result.broker.get_value_history()]\n",
    "        returns = np.array(portfolio_value) / portfolio_value[0] - 1\n",
    "        all_returns.append(returns)\n",
    "\n",
    "    return all_returns\n",
    "\n",
    "# 可视化函数\n",
    "def visualize(selected_stocks, all_returns, time_scale='D'):\n",
    "    plt.figure(figsize=(16, 12))\n",
    "\n",
    "    # 绘制收益率曲线\n",
    "    plt.subplot(2, 1, 1)\n",
    "    for i, stock in enumerate(selected_stocks):\n",
    "        if time_scale == 'D':\n",
    "            dates = pd.date_range(start='2020-01-01', end='2023-12-31')\n",
    "        elif time_scale == 'W':\n",
    "            dates = pd.date_range(start='2020-01-01', end='2023-12-31', freq='W')\n",
    "        elif time_scale == '1min':\n",
    "            dates = pd.date_range(start='2023-01-01', end='2023-12-31', freq='1min')\n",
    "        plt.plot(dates, all_returns[i], label=stock)\n",
    "    plt.xlabel('Date')\n",
    "    plt.ylabel('Return')\n",
    "    plt.title('Returns Curve of Different Stocks')\n",
    "    plt.legend()\n",
    "\n",
    "    # 选择一只股票绘制K线和交易信号\n",
    "    sample_stock = selected_stocks[0]\n",
    "    if time_scale == 'D':\n",
    "        df = pro.daily(ts_code=sample_stock, start_date='20200101', end_date='20231231')\n",
    "    elif time_scale == 'W':\n",
    "        df = pro.weekly(ts_code=sample_stock, start_date='20200101', end_date='20231231')\n",
    "    elif time_scale == '1min':\n",
    "        df = pro.mins(ts_code=sample_stock, freq='1min', start_date='20230101', end_date='20231231')\n",
    "\n",
    "    df['trade_date'] = pd.to_datetime(df['trade_date'])\n",
    "    df.set_index('trade_date', inplace=True)\n",
    "\n",
    "    plt.subplot(2, 1, 2)\n",
    "    plt.plot(df['close'], label='Close Price')\n",
    "    buy_signals = []\n",
    "    sell_signals = []\n",
    "    for i in range(len(df)):\n",
    "        if i > 0:\n",
    "            if df['close'].iloc[i] > df['close'].iloc[i - 1] and df['close'].iloc[i - 1] < df['close'].iloc[i - 2]:\n",
    "                buy_signals.append(df.index[i])\n",
    "            elif df['close'].iloc[i] < df['close'].iloc[i - 1] and df['close'].iloc[i - 1] > df['close'].iloc[i - 2]:\n",
    "                sell_signals.append(df.index[i])\n",
    "    plt.scatter(buy_signals, df['close'].loc[buy_signals], color='green', marker='^', label='Buy Signal')\n",
    "    plt.scatter(sell_signals, df['close'].loc[sell_signals], color='red', marker='v', label='Sell Signal')\n",
    "    plt.xlabel('Date')\n",
    "    plt.ylabel('Price')\n",
    "    plt.title(f'K - line and Trading Signals of {sample_stock}')\n",
    "    plt.legend()\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    selected_stocks = select_and_sort_stocks()\n",
    "    all_returns = backtest(selected_stocks, time_scale='D')\n",
    "    visualize(selected_stocks, all_returns, time_scale='D')\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8a5296f",
   "metadata": {},
   "outputs": [],
   "source": [
    "conda install -c conda-forge backtrader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "905e3092",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
